package com.example.sauc;

import com.google.gson.Gson;
import com.google.gson.JsonObject;
import java.io.*;
import java.util.*;
import java.util.concurrent.*;
import java.util.zip.*;
import javax.sound.sampled.*;
import okhttp3.*;
import okhttp3.logging.HttpLoggingInterceptor;
import okhttp3.logging.HttpLoggingInterceptor.Level;
import okio.ByteString;
import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import java.nio.charset.StandardCharsets;
import java.time.Instant;

public class SaucWebSocketDemo {
    // 协议常量
    private static final String WS_URL = "wss://openspeech-ws.volcengineapi.com/api/v1/asr";
    private static final String SIGNATURE_METHOD = "HmacSHA256";
    private static final String SIGNATURE_VERSION = "1.0";

    // 线程控制
    private static volatile boolean isRunning = true;
    private static final BlockingQueue<byte[]> responseQueue = new LinkedBlockingQueue<>();
    private static final ExecutorService executor = Executors.newFixedThreadPool(2);

    public static void main(String[] args) throws Exception {
        if (args.length != 4) {
            System.err.println("需要传入 4 个参数，依次为 url、appId、token 和 audioFilePath");
            System.exit(1);
        }

        final String url = WS_URL;
        final String appId = args[1];
        final String token = args[2];
        final String audioFilePath = args[3];

        System.out.println("===> 使用参数:");
        System.out.println("URL: " + url);
        System.out.println("AppID: " + appId);
        System.out.println("Token: " + token);
        System.out.println("音频文件: " + audioFilePath);

        // 生成认证信息
        String requestId = UUID.randomUUID().toString();
        long timestamp = Instant.now().getEpochSecond();
        String signature = generateSignature(token, timestamp, requestId);

        // 创建WebSocket客户端
        OkHttpClient client = new OkHttpClient.Builder()
                .pingInterval(60, TimeUnit.SECONDS)
                .readTimeout(60, TimeUnit.SECONDS)
                .writeTimeout(60, TimeUnit.SECONDS)
                .callTimeout(60, TimeUnit.SECONDS)
                .build();

        String queryString = String.format("?app_id=%s&request_id=%s&timestamp=%d", appId, requestId, timestamp);
        String fullUrl = url + queryString;
        System.out.println("===> 完整URL: " + fullUrl);

        Request request = new Request.Builder()
                .url(fullUrl)
                .header("Host", "openspeech-ws.volcengineapi.com")
                .header("Authorization", token)
                .header("Content-Type", "application/json")
                .header("X-Signature", signature)
                .build();

        System.out.println("===> 请求头:");
        Headers headers = request.headers();
        for (int i = 0; i < headers.size(); i++) {
            System.out.println(headers.name(i) + ": " + headers.value(i));
        }

        WebSocket webSocket = client.newWebSocket(request, new WebSocketListener() {
            @Override
            public void onOpen(WebSocket webSocket, Response response) {
                System.out.println("===> 连接已建立");
                // 发送配置信息
                JsonObject config = new JsonObject();
                config.addProperty("app_id", appId);
                config.addProperty("user_id", "test_user");
                config.addProperty("audio_format", "raw");
                config.addProperty("sample_rate", 16000);
                config.addProperty("encode_type", "raw");
                config.addProperty("channel", 1);
                config.addProperty("language", "zh");
                config.addProperty("enable_punctuation", true);
                config.addProperty("enable_interim_result", true);
                config.addProperty("enable_itn", true);
                String configStr = new Gson().toJson(config);
                System.out.println("===> 发送配置: " + configStr);
                webSocket.send(configStr);
            }

            @Override
            public void onMessage(WebSocket webSocket, String text) {
                System.out.println("===> 收到文本消息: " + text);
            }

            @Override
            public void onMessage(WebSocket webSocket, ByteString bytes) {
                responseQueue.offer(bytes.toByteArray());
            }

            @Override
            public void onClosing(WebSocket webSocket, int code, String reason) {
                System.out.println("===> 连接正在关闭: code=" + code + ", reason=" + reason);
                isRunning = false;
            }

            @Override
            public void onFailure(WebSocket webSocket, Throwable t, Response response) {
                System.err.println("===> 连接失败: " + t.getMessage());
                if (response != null) {
                    System.err.println("===> 响应码: " + response.code());
                    try {
                        System.err.println("===> 响应体: " + response.body().string());
                    } catch (IOException e) {
                        e.printStackTrace();
                    }
                    Headers headers = response.headers();
                    System.err.println("===> 响应头:");
                    for (int i = 0; i < headers.size(); i++) {
                        System.err.println(headers.name(i) + ": " + headers.value(i));
                    }
                }
                isRunning = false;
            }
        });

        // 解析音频文件
        try (AudioInputStream ins = AudioSystem.getAudioInputStream(new File(audioFilePath))) {
            AudioFormat format = ins.getFormat();
            System.out.println("音频格式: " + format);
            CountDownLatch latch = new CountDownLatch(2);

            // 启动发送线程
            executor.submit(() -> {
                try {
                    sendAudioData(webSocket, ins);
                } catch (Exception e) {
                    System.err.println("发送线程出错: " + e.getMessage());
                    e.printStackTrace();
                } finally {
                    latch.countDown();
                }
            });

            // 启动接收线程
            executor.submit(() -> {
                try {
                    processResponses();
                } catch (Exception e) {
                    System.err.println("接收线程出错: " + e.getMessage());
                    e.printStackTrace();
                } finally {
                    latch.countDown();
                }
            });

            latch.await();
        }

        executor.shutdown();
        executor.awaitTermination(5, TimeUnit.SECONDS);
    }

    private static void sendAudioData(WebSocket webSocket, AudioInputStream ins) throws IOException {
        byte[] buffer = new byte[3200]; // 每次发送 100ms 的音频
        int bytesRead;
        
        while ((bytesRead = ins.read(buffer)) != -1) {
            if (bytesRead > 0) {
                webSocket.send(ByteString.of(buffer, 0, bytesRead));
            }
            try {
                Thread.sleep(100); // 模拟实时音频流
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                break;
            }
        }
        
        // 发送结束标记
        JsonObject end = new JsonObject();
        end.addProperty("end", true);
        webSocket.send(end.toString());
    }

    private static void processResponses() {
        while (isRunning) {
            try {
                byte[] response = responseQueue.poll(1, TimeUnit.SECONDS);
                if (response != null) {
                    String text = new String(response, StandardCharsets.UTF_8);
                    System.out.println("识别结果: " + text);
                }
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                break;
            }
        }
    }

    private static String generateSignature(String token, long timestamp, String requestId) throws Exception {
        String stringToSign = String.format("app_id=%s&request_id=%s&timestamp=%d", "1322506966", requestId, timestamp);
        System.out.println("===> 待签名字符串: " + stringToSign);
        Mac mac = Mac.getInstance(SIGNATURE_METHOD);
        SecretKeySpec secretKeySpec = new SecretKeySpec(token.getBytes(StandardCharsets.UTF_8), SIGNATURE_METHOD);
        mac.init(secretKeySpec);
        byte[] hash = mac.doFinal(stringToSign.getBytes(StandardCharsets.UTF_8));
        String signature = Base64.getEncoder().encodeToString(hash);
        System.out.println("===> 生成的签名: " + signature);
        return signature;
    }
}